import requests
import threading
import time

# =========================
# 配置
# =========================
URL = "http://127.0.0.1:5000/api/v1/translate"  # Flask 翻译接口
DATA = {
    "text": "一定会成功的",
    "source_lang": "zh-cn",
    "target_lang": "ja"
}  # 请求数据
THREAD_COUNT = 5   # 并发线程数
REQUEST_COUNT = 100  # 总请求数

# =========================
# 统计
# =========================
lock = threading.Lock()
success_count = 0
fail_count = 0

# =========================
# 请求函数
# =========================
def send_request():
    global success_count, fail_count
    try:
        response = requests.post(URL, json=DATA, timeout=5)
        if response.status_code == 200:
            with lock:
                success_count += 1
        else:
            with lock:
                fail_count += 1
    except Exception:
        with lock:
            fail_count += 1

# =========================
# 多线程压测
# =========================
def run_stress_test():
    threads = []
    start_time = time.time()

    for _ in range(REQUEST_COUNT):
        t = threading.Thread(target=send_request)
        threads.append(t)
        t.start()

        # 控制线程数不超过 THREAD_COUNT
        while threading.active_count() > THREAD_COUNT:
            time.sleep(0.01)

    # 等待所有线程结束
    for t in threads:
        t.join()

    end_time = time.time()
    print(f"总请求数: {REQUEST_COUNT}")
    print(f"并发数: {THREAD_COUNT}")
    print(f"成功请求数: {success_count}")
    print(f"失败请求数: {fail_count}")
    print(f"总耗时: {end_time - start_time:.2f} 秒")

if __name__ == "__main__":
    run_stress_test()
